$φ$-Balancing for Mixture-of-Experts Training
ϕ-均衡:面向混合专家训练的凸对偶负载均衡框架
🧭 MoE 训练 · 负载均衡 ★ 深度解读
University of Texas at Austin · Northwestern University
利用严格凸势函数与镜面下降,将 MoE 负载均衡从启发式 batch 统计提升为人口级凸优化,仅需 EMA 跟踪路由概率即实现更稳定均衡与专家专业化
📌 工作参考价值(4 点)
- 用 EMA 替代 batch 统计,消除负载均衡偏差:在电商推荐或 Agent 系统的 MoE 层中,可直接将 ST-MoE 损失中的 batch 频率 f_e 替换为 EMA 路由概率经镜面映射 ∇ϕ(m) 的加权,实现全局均衡,避免因小批次造成的不均匀分配,对长尾流量和冷启动专家更友好。
- 负熵势函数是简单有效的默认选择:实验证明负熵作为 ϕ 在各种规模与粒度下均最优,实现简单(L_aux = Σ p_e * (log m_e + 1)),且对 EMA 衰减率鲁棒(推荐 0.6–0.7),可直接替换现有负载均衡损失而无需额外调参。
- 促进专家专业化,适配多域推荐:ϕ-均衡允许路由根据全局分布动态分配专家,使不同专家自然专注于不同商品类目或行为模式,适合电商多场景联合建模(如搜索、推荐、广告),可尝试冻结部分专家实现领域隔离。
- 工程开销极小,易于集成:仅需在每个 MoE 层维护一份 EMA 向量并增加一步 stop_gradient 的 ∇ϕ 计算,额外算力可忽略,适合大规模分布式训练,能够无缝融入现有 DeepSpeed-MoE 等框架。